#include <iostream>
#include <filesystem>
#include <fstream>
#include <string>
#include <vector>
#include <algorithm>
#include <iomanip>
#include <ctime>
#include <sys/stat.h>

namespace fs = std::filesystem;

class TreeGenerator {
private:
    struct Node {
        std::string name;
        bool is_directory;
        int level;
        std::string permissions;
        std::string size;
        std::string modify_time;
        std::string create_time;
        std::vector<Node> children;
    };
    
    std::string indent = "    ";
    std::string branch = "├── ";
    std::string last_branch = "└── ";
    std::string vertical = "│   ";
    
    Node root;
    
    std::string getPermissions(const fs::path& path) {
        std::string perms;
        fs::perms p = fs::status(path).permissions();
        
        perms += (fs::is_directory(path) ? "d" : "-");
        perms += ((p & fs::perms::owner_read) != fs::perms::none ? "r" : "-");
        perms += ((p & fs::perms::owner_write) != fs::perms::none ? "w" : "-");
        perms += ((p & fs::perms::owner_exec) != fs::perms::none ? "x" : "-");
        perms += ((p & fs::perms::group_read) != fs::perms::none ? "r" : "-");
        perms += ((p & fs::perms::group_write) != fs::perms::none ? "w" : "-");
        perms += ((p & fs::perms::group_exec) != fs::perms::none ? "x" : "-");
        perms += ((p & fs::perms::others_read) != fs::perms::none ? "r" : "-");
        perms += ((p & fs::perms::others_write) != fs::perms::none ? "w" : "-");
        perms += ((p & fs::perms::others_exec) != fs::perms::none ? "x" : "-");
        
        return perms;
    }
    
    std::string formatFileSize(uintmax_t size) {
        const char* units[] = {"B", "KB", "MB", "GB", "TB"};
        int unit_index = 0;
        double file_size = static_cast<double>(size);
        
        while (file_size >= 1024.0 && unit_index < 4) {
            file_size /= 1024.0;
            unit_index++;
        }
        
        std::ostringstream oss;
        oss << std::fixed << std::setprecision(1) << file_size << " " << units[unit_index];
        return oss.str();
    }
    
    std::string formatTime(std::time_t time) {
        char buffer[20];
        std::strftime(buffer, sizeof(buffer), "%Y-%m-%d %H:%M:%S", std::localtime(&time));
        return std::string(buffer);
    }
    
    // 兼容C++17的文件时间转换函数
    std::time_t to_time_t(const fs::file_time_type& ft) {
        // 将file_time_type转换为time_t的通用方法
        auto sctp = std::chrono::time_point_cast<std::chrono::system_clock::duration>(
            ft - fs::file_time_type::clock::now() + std::chrono::system_clock::now());
        return std::chrono::system_clock::to_time_t(sctp);
    }
    
    void buildTree(const fs::path& path, Node& node, int level = 0) {
        if (!fs::exists(path)) return;
        
        // 获取目录中的所有条目并排序
        std::vector<fs::directory_entry> entries;
        for (const auto& entry : fs::directory_iterator(path)) {
            entries.push_back(entry);
        }
        
        // 按名称排序
        std::sort(entries.begin(), entries.end(), 
                  [](const fs::directory_entry& a, const fs::directory_entry& b) {
                      return a.path().filename().string() < b.path().filename().string();
                  });
        
        for (size_t i = 0; i < entries.size(); ++i) {
            const auto& entry = entries[i];
            Node child;
            child.name = entry.path().filename().string();
            child.is_directory = fs::is_directory(entry.status());
            child.level = level + 1;
            
            // 获取文件属性
            child.permissions = getPermissions(entry.path());
            
            if (!child.is_directory) {
                try {
                    child.size = formatFileSize(fs::file_size(entry.path()));
                } catch (const fs::filesystem_error&) {
                    child.size = "Unknown";
                }
            } else {
                child.size = "-";
            }
            
            // 获取修改时间 (兼容C++17)
            auto modify_time = fs::last_write_time(entry.path());
            auto modify_time_t = to_time_t(modify_time);
            child.modify_time = formatTime(modify_time_t);
            
            // 获取创建时间（如果可用）
            try {
                struct stat stat_buf;
                if (stat(entry.path().c_str(), &stat_buf) == 0) {
                    child.create_time = formatTime(stat_buf.st_ctime);
                } else {
                    child.create_time = child.modify_time; // 如果无法获取，使用修改时间
                }
            } catch (...) {
                child.create_time = child.modify_time;
            }
            
            node.children.push_back(child);
            
            if (child.is_directory) {
                buildTree(entry.path(), node.children.back(), level + 1);
            }
        }
    }
    
    void printTree(const Node& node, const std::string& prefix = "", bool is_last = true) {
        std::string current_prefix = prefix + (is_last ? last_branch : branch);
        std::cout << prefix << (is_last ? last_branch : branch) << node.name << std::endl;
        
        for (size_t i = 0; i < node.children.size(); ++i) {
            bool child_is_last = (i == node.children.size() - 1);
            std::string new_prefix = prefix + (is_last ? indent : vertical);
            printTree(node.children[i], new_prefix, child_is_last);
        }
    }
    
    void generateTXT(const Node& node, std::ofstream& file, const std::string& prefix = "", bool is_last = true) {
        std::string current_prefix = prefix + (is_last ? last_branch : branch);
        
        // 输出节点信息
        file << current_prefix << node.name;
        if (!node.is_directory) {
            file << " (" << node.size << ")";
        }
        file << std::endl;
        
        // 输出属性信息（缩进对齐）
        std::string attr_prefix = prefix + (is_last ? "    " : vertical + "    ");
        file << attr_prefix << "权限: " << node.permissions << std::endl;
        file << attr_prefix << "修改时间: " << node.modify_time << std::endl;
        file << attr_prefix << "创建时间: " << node.create_time << std::endl;
        
        // 递归处理子节点
        for (size_t i = 0; i < node.children.size(); ++i) {
            bool child_is_last = (i == node.children.size() - 1);
            std::string new_prefix = prefix + (is_last ? indent : vertical);
            generateTXT(node.children[i], file, new_prefix, child_is_last);
        }
    }
    
    void generateHTML(const Node& node, std::ofstream& file, bool is_root = false) {
        if (is_root) {
            file << "<!DOCTYPE html>\n";
            file << "<html>\n";
            file << "<head>\n";
            file << "    <meta charset=\"UTF-8\">\n";
            file << "    <title>Directory Tree</title>\n";
            file << "    <style>\n";
            file << "        body { font-family: Arial, sans-serif; margin: 20px; background-color: #f5f5f5; }\n";
            file << "        .container { background: white; padding: 20px; border-radius: 5px; box-shadow: 0 2px 5px rgba(0,0,0,0.1); }\n";
            file << "        .tree { list-style-type: none; padding-left: 0; }\n";
            file << "        .tree ul { list-style-type: none; padding-left: 20px; }\n";
            file << "        .folder { color: #1e90ff; font-weight: bold; cursor: pointer; }\n";
            file << "        .file { color: #333; }\n";
            file << "        .toggle { cursor: pointer; margin-right: 5px; display: inline-block; width: 12px; }\n";
            file << "        .hidden { display: none; }\n";
            file << "        .attributes { font-size: 0.85em; color: #666; margin-left: 20px; }\n";
            file << "        .attr-line { margin: 2px 0; }\n";
            file << "        h2 { color: #333; border-bottom: 2px solid #1e90ff; padding-bottom: 10px; }\n";
            file << "        button { background: #1e90ff; color: white; border: none; padding: 5px 10px; border-radius: 3px; cursor: pointer; }\n";
            file << "        button:hover { background: #0077e6; }\n";
            file << "        .file-size { color: #888; font-style: italic; margin-left: 10px; }\n";
            file << "    </style>\n";
            file << "    <script>\n";
            file << "        function toggleFolder(element) {\n";
            file << "            const parent = element.parentElement;\n";
            file << "            const children = parent.querySelector('ul');\n";
            file << "            const attributes = parent.querySelector('.attributes');\n";
            file << "            if (children) {\n";
            file << "                children.classList.toggle('hidden');\n";
            file << "                if (attributes) {\n";
            file << "                    attributes.classList.toggle('hidden');\n";
                    file << "                }\n";
            file << "                element.textContent = children.classList.contains('hidden') ? '▶' : '▼';\n";
            file << "            }\n";
            file << "        }\n";
            file << "        function expandAll() {\n";
            file << "            document.querySelectorAll('.tree ul').forEach(ul => ul.classList.remove('hidden'));\n";
            file << "            document.querySelectorAll('.attributes').forEach(attr => attr.classList.remove('hidden'));\n";
            file << "            document.querySelectorAll('.toggle').forEach(toggle => toggle.textContent = '▼');\n";
            file << "        }\n";
            file << "        function collapseAll() {\n";
            file << "            document.querySelectorAll('.tree ul').forEach(ul => ul.classList.add('hidden'));\n";
            file << "            document.querySelectorAll('.attributes').forEach(attr => attr.classList.add('hidden'));\n";
            file << "            document.querySelectorAll('.toggle').forEach(toggle => toggle.textContent = '▶');\n";
            file << "            const rootUl = document.querySelector('.tree > ul');\n";
            file << "            if (rootUl) rootUl.classList.remove('hidden');\n";
            file << "            const rootToggle = document.querySelector('.tree > .folder > .toggle');\n";
            file << "            if (rootToggle) rootToggle.textContent = '▼';\n";
            file << "        }\n";
            file << "    </script>\n";
            file << "</head>\n";
            file << "<body>\n";
            file << "<div class=\"container\">\n";
            file << "<h2>Directory Tree: " << fs::current_path().string() << "</h2>\n";
            file << "<div style=\"margin-bottom: 15px;\">\n";
            file << "    <button onclick=\"expandAll()\">Expand All</button>\n";
            file << "    <button onclick=\"collapseAll()\" style=\"margin-left: 10px;\">Collapse All</button>\n";
            file << "</div>\n";
            file << "<div class=\"tree\">\n";
        }
        
        // 输出当前节点
        if (node.is_directory) {
            file << "<div class=\"folder\">";
            if (!is_root) {
                file << "<span class=\"toggle\" onclick=\"toggleFolder(this)\">▼</span> ";
            }
            file << node.name << "\n";
            
            // 输出属性信息
            file << "<div class=\"attributes\">\n";
            file << "    <div class=\"attr-line\">权限: " << node.permissions << "</div>\n";
            file << "    <div class=\"attr-line\">修改时间: " << node.modify_time << "</div>\n";
            file << "    <div class=\"attr-line\">创建时间: " << node.create_time << "</div>\n";
            file << "</div>\n";
            
            file << "<ul" << (is_root ? "" : " class=\"hidden\"") << ">\n";
            
            for (const auto& child : node.children) {
                file << "<li>";
                generateHTML(child, file, false);
                file << "</li>\n";
            }
            
            file << "</ul>\n";
            file << "</div>\n";
        } else {
            file << "<span class=\"file\">" << node.name << "<span class=\"file-size\">(" << node.size << ")</span></span>\n";
            
            // 输出属性信息
            file << "<div class=\"attributes\">\n";
            file << "    <div class=\"attr-line\">权限: " << node.permissions << "</div>\n";
            file << "    <div class=\"attr-line\">修改时间: " << node.modify_time << "</div>\n";
            file << "    <div class=\"attr-line\">创建时间: " << node.create_time << "</div>\n";
            file << "</div>\n";
        }
        
        if (is_root) {
            file << "</div>\n"; // 关闭tree div
            file << "</div>\n"; // 关闭container div
            file << "</body>\n";
            file << "</html>\n";
        }
    }

public:
    void generate(const std::string& path = ".", const std::string& output_file = "") {
        // 构建根节点
        root.name = fs::path(path).filename().string();
        root.is_directory = true;
        root.level = 0;
        root.permissions = getPermissions(path);
        root.size = "-";
        
        // 获取根目录的时间信息 (兼容C++17)
        auto modify_time = fs::last_write_time(path);
        auto modify_time_t = to_time_t(modify_time);
        root.modify_time = formatTime(modify_time_t);
        
        try {
            struct stat stat_buf;
            if (stat(path.c_str(), &stat_buf) == 0) {
                root.create_time = formatTime(stat_buf.st_ctime);
            } else {
                root.create_time = root.modify_time;
            }
        } catch (...) {
            root.create_time = root.modify_time;
        }
        
        // 构建树结构
        buildTree(path, root);
        
        if (output_file.empty()) {
            // 输出到终端
            std::cout << "." << std::endl;
            for (size_t i = 0; i < root.children.size(); ++i) {
                bool is_last = (i == root.children.size() - 1);
                printTree(root.children[i], "", is_last);
            }
        } else {
            // 根据文件扩展名判断输出格式
            std::string ext = output_file.substr(output_file.find_last_of(".") + 1);
            
            if (ext == "html" || ext == "htm") {
                // 输出到HTML文件
                std::ofstream file(output_file);
                if (!file.is_open()) {
                    std::cerr << "Error: Cannot create file " << output_file << std::endl;
                    return;
                }
                
                generateHTML(root, file, true);
                file.close();
                std::cout << "HTML file generated: " << output_file << std::endl;
            } else if (ext == "txt") {
                // 输出到TXT文件
                std::ofstream file(output_file);
                if (!file.is_open()) {
                    std::cerr << "Error: Cannot create file " << output_file << std::endl;
                    return;
                }
                
                file << "Directory Tree: " << fs::current_path().string() << std::endl;
                file << "Generated on: " << formatTime(std::time(nullptr)) << std::endl;
                file << "==========================================" << std::endl;
                
                file << "." << std::endl;
                file << "权限: " << root.permissions << std::endl;
                file << "修改时间: " << root.modify_time << std::endl;
                file << "创建时间: " << root.create_time << std::endl;
                
                for (size_t i = 0; i < root.children.size(); ++i) {
                    bool is_last = (i == root.children.size() - 1);
                    generateTXT(root.children[i], file, "", is_last);
                }
                
                file.close();
                std::cout << "TXT file generated: " << output_file << std::endl;
            } else {
                std::cerr << "Error: Unsupported output format. Use .html, .htm or .txt" << std::endl;
            }
        }
    }
};

void printUsage() {
    std::cout << "Usage:\n";
    std::cout << "  tree                    - Display directory tree in terminal\n";
    std::cout << "  tree . -o tree.html     - Generate directory tree as HTML file\n";
    std::cout << "  tree . -o tree.txt      - Generate directory tree as TXT file\n";
}

int main(int argc, char* argv[]) {
    TreeGenerator tree;
    
    if (argc == 1) {
        // tree
        tree.generate();
    } else if (argc == 4) {
        // tree . -o output.file
        std::string path = argv[1];
        std::string option = argv[2];
        std::string output_file = argv[3];
        
        if (path == "." && option == "-o") {
            tree.generate(".", output_file);
        } else {
            printUsage();
            return 1;
        }
    } else {
        printUsage();
        return 1;
    }
    
    return 0;
}
